import soundfile as sf
import numpy as np
import sys

#一个噪声，wav叠加的功能

def add_noise(noisewav,cleanwav,snr):
    # read noisewav
    noise, fs = sf.read(noisewav)
    print(noisewav,'采样率',fs)
    #read cleanwav
    clean, Fs = sf.read(cleanwav)
    print(cleanwav,'采样率',Fs)
    # add noise
    if fs == Fs and len(clean) <= len(noise):
        # 纯净语音能量
        cleanenergy = np.sum(np.power(clean,2))
        # 随机索引与clean长度相同的noise信号
        ind = np.random.randint(1, len(noise) - len(clean) + 1)
        noiselen = noise[ind:len(clean) + ind]
        # 噪声语音能量
        noiseenergy = np.sum(np.power(noiselen,2))
        # 噪声等级系数
        noiseratio = np.sqrt((cleanenergy / noiseenergy) / (np.power(10, snr * 0.1)))
        # 随机索引与clean长度相同的noise信号
        noisyAudio = clean + noise[ind:len(clean)+ind] * noiseratio
        # write wav
        noisywavname = cleanwav[:-4] + "_"+noisewav[:-4] + "_snr"+str(snr)+".wav"
        sf.write(noisywavname, noisyAudio, 16000)

    elif fs == Fs and len(clean) >= len(noise):
        # 纯净语音能量
        cleanenergy = np.sum(np.power(clean, 2))
        # 随机索引与clean长度相同的noise信号
        ind = np.random.randint(1, len(clean) - len(noise) + 1)
        noiselen = clean[ind:len(noise) + ind]
        # 噪声语音能量
        noiseenergy = np.sum(np.power(noiselen, 2))
        # 噪声等级系数
        noiseratio = np.sqrt((cleanenergy / noiseenergy) / (np.power(10, snr * 0.1)))
        # 随机索引与clean长度相同的noise信号
        noisyAudio = noise + clean[ind:len(noise) + ind] * noiseratio
        # write wav
        noisywavname = cleanwav[:-4] + "_" + noisewav[:-4] + "_snr" + str(snr) + ".wav"
        sf.write(noisywavname, noisyAudio, 16000)
    else:
        print("与采样率不一样")
        sys.exit(-1)


if __name__ == '__main__':
    noisewav = "5_1894_20170630160624.wav"
    cleanwav = "5_1894_20170809152805.wav"
    snr = 5
    add_noise(noisewav, cleanwav, snr)
    